import matplotlib.pyplot as plt

from PIL import Image


# ============================================= 图片处理相关 ========================================================
# 保存图片
def draw_img(data):
    plt.figure()
    plt.imshow(data)
    plt.axis('off')  # 这将隐藏x轴和y轴
    # plt.colorbar()
    # plt.title('mag_E')
    # plt.xlabel('Theta')
    # plt.ylabel('Phi')
    plt.show()

def save_img(path_img, data):
    plt.figure()
    plt.imshow(data)
    plt.axis('off')  # 这将隐藏x轴和y轴
    # plt.colorbar()
    # plt.title('mag_E')
    # plt.xlabel('Theta')
    # plt.ylabel('Phi')
    # plt.show()
    # plt.savefig(path_img, dpi=300, bbox_inches='tight', pad_inches=0)
    plt.savefig(path_img, dpi=200, bbox_inches='tight', pad_inches=0)
    plt.close(fig=None)  # 关闭当前图像窗口，如果fig=None，则默认关闭最近打开的图像窗口


def draw_img_xyz(data, x, y):
    # 绘制图像
    plt.figure(figsize=(10, 8))
    plt.pcolormesh(x, y, data, shading='auto', cmap='plasma')
    # plt.pcolormesh(x, y, np.abs(pattern_xyz), shading='auto', cmap='viridis')
    # plt.colorbar(label='Pattern Magnitude')
    plt.axis('equal')
    plt.axis('off')  # 关闭坐标轴显示
    plt.show()
    # plt.savefig(path_img, dpi=200, bbox_inches='tight', pad_inches=0)
    # plt.close(fig=None)  # 关闭当前图像窗口，如果fig=None，则默认关闭最近打开的图像窗口


def save_img_xyz(path_img, data, x, y):
    # 绘制图像
    plt.figure(figsize=(10, 8))
    plt.pcolormesh(x, y, data, shading='auto', cmap='viridis')
    # plt.pcolormesh(x, y, np.abs(pattern_xyz), shading='auto', cmap='viridis')
    # plt.colorbar(label='Pattern Magnitude')
    plt.axis('equal')
    plt.axis('off')  # 关闭坐标轴显示
    # plt.show()
    plt.savefig(path_img, dpi=200, bbox_inches='tight', pad_inches=0)
    plt.close(fig=None)  # 关闭当前图像窗口，如果fig=None，则默认关闭最近打开的图像窗口


from matplotlib.colors import LinearSegmentedColormap


def save_img_xyz_2(path_img, data, x, y):
    plt.figure(figsize=(10, 8))

    # 自定义高对比度色图
    colors = [(0, 0, 1), (1, 0, 0)]  # 蓝→红渐变
    custom_cmap = LinearSegmentedColormap.from_list("custom", colors, N=256)

    # 绘制图像（结合方法2和方法3）
    plt.pcolormesh(
        x, y, data,
        cmap=custom_cmap,
        vmin=data.min() * 0.4,  # 压缩低值范围
        vmax=data.max() * 0.8
    )

    # 添加颜色条
    # cbar = plt.colorbar(pad=0.02)
    # cbar.set_label("Intensity", fontsize=12)
    # cbar.set_ticks([data.min() * 0.1, data.mean(), data.max() * 0.9])
    # cbar.set_ticklabels(["Low", "Medium", "High"])

    plt.axis('equal')
    plt.axis('off')
    plt.savefig(path_img, dpi=200, bbox_inches='tight', pad_inches=0)
    plt.close()


def draw_img_xyz_2(data, x, y):
    plt.figure(figsize=(10, 8))

    # 自定义高对比度色图
    colors = [(0, 0, 1), (1, 0, 0)]  # 蓝→红渐变
    custom_cmap = LinearSegmentedColormap.from_list("custom", colors, N=256)

    # 绘制图像（结合方法2和方法3）
    plt.pcolormesh(
        x, y, data,
        cmap=custom_cmap,
        vmin=data.min() * 0.4,  # 压缩低值范围
        vmax=data.max() * 0.8
    )

    # 添加颜色条
    # cbar = plt.colorbar(pad=0.02)
    # cbar.set_label("Intensity", fontsize=12)
    # cbar.set_ticks([data.min() * 0.1, data.mean(), data.max() * 0.9])
    # cbar.set_ticklabels(["Low", "Medium", "High"])

    plt.axis('equal')
    plt.axis('off')
    plt.show()


def plot_line_chart(list_data, title, x_label, y_label, legend_label):
    """
    绘制折线图的函数。
    参数:
    list_data (list): 要绘制的列表数据。
    title (str): 图表的标题。
    x_label (str): X轴标签。
    y_label (str): Y轴标签。
    legend_label (str): 图例标签。
    """
    # 创建一个新的绘图对象
    plt.figure()
    # 绘制折线图
    plt.plot(list_data, marker='o', linestyle='-', color='b')
    # 设置标题
    plt.title(title)
    # 设置X轴标签
    plt.xlabel(x_label)
    # 设置Y轴标签
    plt.ylabel(y_label)
    # 添加网格
    plt.grid(True)
    # 显示图例
    plt.legend([legend_label])
    # 显示图形
    plt.show()


def save_line_chart(path_img, list_data, title, x_label, y_label, legend_label):
    """
    保存折线图的函数。
    参数:
    list_data (list): 要绘制的列表数据。
    title (str): 图表的标题。
    x_label (str): X轴标签。
    y_label (str): Y轴标签。
    legend_label (str): 图例标签。
    path_img (str): 保存图片的路径。
    """
    # 创建一个新的绘图对象
    plt.figure()
    # 绘制折线图
    plt.plot(list_data, marker='o', linestyle='-', color='b')
    # 设置标题
    plt.title(title)
    # 设置X轴标签
    plt.xlabel(x_label)
    # 设置Y轴标签
    plt.ylabel(y_label)
    # 添加网格
    plt.grid(True)
    # 显示图例
    plt.legend([legend_label])
    # 保存图形到指定路径
    plt.savefig(path_img)
    # 关闭图形，以释放内存
    plt.close()


def save_concat_images_horizontal_2(image_path_left, image_path_right, output_path):
    # 打开两张图片
    image_left = Image.open(image_path_left)
    image_right = Image.open(image_path_right)
    # 确保两张图片具有相同的高度
    if image_left.height != image_right.height:
        # 选择较小的高度来保持比例
        min_height = min(image_left.height, image_right.height)
        # 调整第一张图片大小
        image_left = image_left.resize((int(image_left.width * (min_height / image_left.height)), min_height),
                                   resample=Image.Resampling.LANCZOS)
        # 调整第二张图片大小
        image_right = image_right.resize((int(image_right.width * (min_height / image_right.height)), min_height),
                                    resample=Image.Resampling.LANCZOS)
    # 创建一个新的空白图像，宽度为两张图片宽度之和，高度与原图相同
    new_image = Image.new('RGB', (image_left.width + image_right.width, image_left.height))
    # 将第一张图片粘贴到新图像的左侧
    new_image.paste(image_left, (0, 0))
    # 将第二张图片粘贴到新图像的右侧
    new_image.paste(image_right, (image_left.width, 0))
    # 保存结果
    new_image.save(output_path)


def save_concat_images_horizontal_3(image_path_left, image_path_middle, image_path_right, output_path):
    # 打开三张图片
    image_left = Image.open(image_path_left)
    image_middle = Image.open(image_path_middle)
    image_right = Image.open(image_path_right)
    # 确保三张图片具有相同的高度
    heights = [image_left.height, image_middle.height, image_right.height]
    min_height = min(heights)
    # 调整图片大小以保持比例
    image_left = image_left.resize((int(image_left.width * (min_height / image_left.height)), min_height),
                                   resample=Image.Resampling.LANCZOS)
    image_middle = image_middle.resize((int(image_middle.width * (min_height / image_middle.height)), min_height),
                                       resample=Image.Resampling.LANCZOS)
    image_right = image_right.resize((int(image_right.width * (min_height / image_right.height)), min_height),
                                     resample=Image.Resampling.LANCZOS)
    # 创建一个新的空白图像，宽度为三张图片宽度之和，高度与原图相同
    new_image = Image.new('RGB', (image_left.width + image_middle.width + image_right.width, min_height))
    # 将第一张图片粘贴到新图像的左侧
    new_image.paste(image_left, (0, 0))
    # 将第二张图片粘贴到新图像的中间
    new_image.paste(image_middle, (image_left.width, 0))
    # 将第三张图片粘贴到新图像的右侧
    new_image.paste(image_right, (image_left.width + image_middle.width, 0))
    # 保存结果
    new_image.save(output_path)


# 画图: 3幅 * 1行
def plot_images_3x1(data1, data2, data3):
    # 创建一个新的图像
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(9, 3))
    # 在左上方绘制 list_data1[i]
    im1 = ax1.imshow(data1, aspect='auto')
    ax1.set_title(f"phaseBit")
    ax1.axis('equal')
    ax1.axis('off')
    # plt.colorbar(im1, ax=ax1)     # 注释掉能关闭图例
    # 在中上方绘制 list_data2[i]
    im2 = ax2.imshow(data2, aspect='auto')
    ax2.set_title(f"phaseBit_new")
    ax2.axis('equal')
    ax2.axis('off')
    # plt.colorbar(im2, ax=ax2)
    # 在右上方绘制 list_data3[i]
    im3 = ax3.imshow(data3, aspect='auto')
    ax3.set_title(f"diff")
    ax3.axis('equal')
    ax3.axis('off')
    # plt.colorbar(im3, ax=ax3)
    # 在下方放置文本
    # fig.text(0.5, 0.01, str_label[3], ha='center', fontsize=12)
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, bottom=0.05)
    # 保存图像
    plt.show()


# 画图: 3幅 * 2行 + 1行文字
def plot_images_3x2(data1, data2, data3, data4, data5, data6, img_text):
    # 创建一个新的图像，包含2行3列的子图
    fig, axes = plt.subplots(2, 3, figsize=(18, 9))
    # 在第一行绘制三张图
    im1 = axes[0, 0].imshow(data1, aspect='auto')
    axes[0, 0].set_title(f"phaseBit(t)")
    axes[0, 0].axis('equal')
    axes[0, 0].axis('off')
    im2 = axes[0, 1].imshow(data2, aspect='auto')
    axes[0, 1].set_title(f"phaseBit-NN(t-1)")
    axes[0, 1].axis('equal')
    axes[0, 1].axis('off')
    im3 = axes[0, 2].imshow(data3, aspect='auto')
    axes[0, 2].set_title(f"phaseBit-NN(t)")
    axes[0, 2].axis('equal')
    axes[0, 2].axis('off')
    # 在第二行绘制三张图
    im4 = axes[1, 0].imshow(data4, aspect='auto')
    axes[1, 0].set_title(f"pattern")
    axes[1, 0].axis('equal')
    axes[1, 0].axis('off')
    im5 = axes[1, 1].imshow(data5, aspect='auto')
    axes[1, 1].set_title(f"diff(phaseBit(t), phaseBit-NN(t))")
    axes[1, 1].axis('equal')
    axes[1, 1].axis('off')
    im6 = axes[1, 2].imshow(data6, aspect='auto')
    axes[1, 2].set_title(f"pattern-NN")
    axes[1, 2].axis('equal')
    axes[1, 2].axis('off')
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, hspace=0.1, bottom=0.05)
    # 在下方放置文本
    fig.text(0.5, 0.01, img_text, ha='center', fontsize=12)
    # 显示图像
    plt.show()


# 保存图: 3幅 * 2行 + 1行文字
def save_images_3x2(path_img,
                    data1, text1, data2, text2, data3, text3,
                    data4, text4, data5, text5, data6, text6,
                    img_text):
    # 创建一个新的图像，包含2行3列的子图
    fig, axes = plt.subplots(2, 3, figsize=(18, 9))
    # 在第一行绘制三张图
    im1 = axes[0, 0].imshow(data1, aspect='auto')
    axes[0, 0].set_title(text1)
    axes[0, 0].axis('equal')
    axes[0, 0].axis('off')
    im2 = axes[0, 1].imshow(data2, aspect='auto')
    axes[0, 1].set_title(text2)
    axes[0, 1].axis('equal')
    axes[0, 1].axis('off')
    im3 = axes[0, 2].imshow(data3, aspect='auto')
    axes[0, 2].set_title(text3)
    axes[0, 2].axis('equal')
    axes[0, 2].axis('off')
    # 在第二行绘制三张图
    im4 = axes[1, 0].imshow(data4, aspect='auto')
    axes[1, 0].set_title(text4)
    axes[1, 0].axis('equal')
    axes[1, 0].axis('off')
    im5 = axes[1, 1].imshow(data5, aspect='auto')
    axes[1, 1].set_title(text5)
    axes[1, 1].axis('equal')
    axes[1, 1].axis('off')
    im6 = axes[1, 2].imshow(data6, aspect='auto')
    axes[1, 2].set_title(text6)
    axes[1, 2].axis('equal')
    axes[1, 2].axis('off')
    # 调整子图间距
    plt.subplots_adjust(wspace=0.1, hspace=0.1, bottom=0.05)
    # 在下方放置文本
    fig.text(0.5, 0.01, img_text, ha='center', fontsize=12)
    # 显示图像
    # plt.show()
    # 保存图像
    plt.savefig(path_img)
    # 关闭当前图像，释放内存
    plt.close(fig)




if __name__ == '__main__':
    # 使用示例
    save_concat_images_horizontal_2('path/to/first/image.jpg', 'path/to/second/image.jpg', 'path/to/save/new_image.jpg')